#coding: utf-8
import os
import sys
import json
import codecs
from collections import OrderedDict

def reader(input):
    with open(input, 'r') as fd:
        for line in fd:
            yield json.loads(line)

def parse_yesno(sample):
    question = sample['question']
    if '是否' in question:
        no = False
        for doc in sample['documents']:
            for para in doc['paragraphs']:
                if '否' in para or '不' in para:
                    no = True
        sample['pred_answers'] = ['否'] if no else ['是']
    return sample

def output(pred, output_file):
    with codecs.open(output_file, 'w', encoding='utf8') as fd:
        for q_id, data in pred.items():
            ans = data['pred_answers'][0].replace('\t', '').replace(' ', '')
            fd.write(f'{q_id}\t{ans}\n')

def merge_meta(meta_file, pred_file):
    meta_dict = OrderedDict()
    for data in reader(meta_file):
        q_id = data['question_id']
        meta_dict[q_id] = data
    pred_dict = OrderedDict()
    for data in reader(pred_file):
        q_id = data['question_id']
        data['pred_answers'] = data['answers']
        pred_dict[q_id] = data
        pred_dict[q_id].update(meta_dict[q_id])
        pred_dict[q_id] = parse_yesno(pred_dict[q_id])
    return pred_dict

def block_format(sample):
    line = 'question_id:%s\n' % sample['question_id']
    line += 'question_type:%s\n' % sample['question_type']
    line += 'question:%s\n' % sample['question']
    for doc in sample['documents']:
        line += 'title:%s\n' % doc['title']
        line += 'paragraph:%s\n' % '+'.join(doc['paragraphs'])
    for key in ['yesno_answers', 'entity_answers', 'answers', 'scores']:
        temp = sample.get(key, ['--'])[0]
        line += f'{key}:{temp}\n'
    line += 'pred_answers:%s' % sample['pred_answers'][0] + '\n'
    line += '----' * 40
    return line

def compare(pred_dict):
    with open('cmp.txt', 'w') as fd:
        qids = sorted(pred_dict.keys())
        for q_id in qids:
            data = pred_dict[q_id]
            for doc in data['documents']:
                del doc['segmented_paragraphs']
                del doc['segmented_title']
            #del data['documents']
            #del data['question']
            del data['segmented_question']
            del data['segmented_answers']
            if data['answers'] == []:
                del data['answers']
            if data['yesno_answers'] == []:
                del data['yesno_answers']
            if data['entity_answers'] == [[]]:
                del data['entity_answers']
            if data['question_type'] == 'entity':
                #fd.write(json.dumps(data, ensure_ascii=False) + '\n')
                fd.write(block_format(data) + '\n')

if __name__ == '__main__':
    meta_file = sys.argv[1]
    pred_file = sys.argv[2]
    output_file = sys.argv[3]
    pred_dict = merge_meta(meta_file, pred_file)
    output(pred_dict, output_file)
    compare(pred_dict)
    print('done')
